##---------------------------------------------------------------
## Required libraries for computing
##---------------------------------------------------------------

#----- Parallel computing
library(doParallel)

#----- Make clusters based on the number of CPU cores
cl <- makeCluster(detectCores())
registerDoParallel(cl)
getDoParWorkers()

#----- Support parallel excution
library(foreach)

#----- Stan (HMC) for R
library(rstan)

##---------------------------------------------------------------
## Extract MCMC samples from Stan outputs
##---------------------------------------------------------------

#----- Load MCMC samples and Data
load("MCMCsample.RData")
load("Master.RData")
load("SurfaceFrame1.RData")
#------ Set Treatment (TRT), Outcome (OUT), Mediators (M) and Covariates (X)
Data <- Master
OUT <- Data$PM.2.5
TRT <- Data$SO2.SC
M <- cbind(Data$SO2_Annual, Data$NOx_Annual, Data$CO2_Annual)
X <- cbind(Data$S_n_CR, Data$NumNOxControls, Data$Heat_Input/100000, Data$Barometric_Pressure, Data$Temperature,  Data$PctCapacity, Data$sulfur_Content, Data$Phase2_Indicator, Data$Operating_Time/1000)

dim.cov <- dim(X)[2] #<--------- Num. of Covariates

#------ Variables by treatments
x0 <- X[which(TRT==0),]
x1 <- X[which(TRT==1),]

y0 <- OUT[which(TRT==0)]
y1 <- OUT[which(TRT==1)]

m0 <- log(M[which(TRT==0),])
m1 <- log(M[which(TRT==1),])

n0 <- dim(x0)[1]
n1 <- dim(x1)[1]

#----- Extract MCMC samples from each marginal distribution

data.y0 <- extract(fit.y0)
data.y1 <- extract(fit.y1)
data.m10 <- extract(fit.m1_0)
data.m11 <- extract(fit.m1_1)
data.m20 <- extract(fit.m2_0)
data.m21 <- extract(fit.m2_1)
data.m30 <- extract(fit.m3_0)
data.m31 <- extract(fit.m3_1)

#----- Parameters of the models under Z=1

gamma10 <- data.y1$gamma0
gamma11 <- data.y1$gamma1
w1 <- data.y1$W
beta1_10 <- data.m11$gamma0
beta1_11 <- data.m11$gamma1
beta2_10 <- data.m21$gamma0
beta2_11 <- data.m21$gamma1
beta3_10 <- data.m31$gamma0
beta3_11 <- data.m31$gamma1
s1_1 <- data.m11$W
s2_1 <- data.m21$W
s3_1 <- data.m31$W
psi1 <- data.y1$psi
psi1_1 <- data.m11$psi
psi2_1 <- data.m21$psi
psi3_1 <- data.m31$psi


#----- Parameters of the models under Z=0

gamma00 <- data.y0$gamma0
gamma01 <- data.y0$gamma1
w0 <- data.y0$W
beta1_00 <- data.m10$gamma0
beta1_01 <- data.m10$gamma1
beta2_00 <- data.m20$gamma0
beta2_01 <- data.m20$gamma1
beta3_00 <- data.m30$gamma0
beta3_01 <- data.m30$gamma1
s1_0 <- data.m10$W
s2_0 <- data.m20$W
s3_0 <- data.m30$W
psi0 <- data.y0$psi
psi1_0 <- data.m10$psi
psi2_0 <- data.m20$psi
psi3_0 <- data.m30$psi


#----- Setting Thinning and Burn-in's
Thin <- 5
Burn0 <- 18000  # Extra burn-in periods for Models under Z=0
Burn1 <- 18000   # Extra burn-in periods for Models under Z=1

#----- Set the number of post-processing steps, N
n.iter <- 200


##---------------------------------------------------------------
## 'Main' function on each cluster (parallel)
##---------------------------------------------------------------

main<-function(temp){


    #----- Require to fit multivariate normal distributions
    library(mnormt)

    joint0 <- cbind(y0, m0);  joint1 <- cbind(y1, m1)
    #----- Estimating condtional correlations
    fit0 <- lm(joint0 ~ x0)
    fit1 <- lm(joint1 ~ x1)

    cov0 <- 1 / n0 * t(joint0 - predict(fit0, newdata=data.frame(x0))) %*% (joint0-predict(fit0, newdata = data.frame(x0)))
    cov1 <- 1 / n1 * t(joint1 - predict(fit1, newdata=data.frame(x1))) %*% (joint1-predict(fit1, newdata = data.frame(x1)))

    #----- Under Normality, the following relatioship between
    #----- Spearman and Pearson correlations holds (Kruskal, 1958)
    cor0 <- 6 / pi * asin(matrix(apply(expand.grid(1:4,1:4), 1, function(x) cov0[x[1],x[2]] / sqrt(cov0[x[1],x[1]] * cov0[x[2],x[2]])),4,4)/2)
    cor1 <- 6 / pi * asin(matrix(apply(expand.grid(1:4,1:4), 1, function(x) cov1[x[1],x[2]] / sqrt(cov1[x[1],x[1]]*cov1[x[2],x[2]])),4,4)/2)


    #----- Sensitivity parameters
    sens1 <- function(j,k) 2 * min(abs(cor0[j+1,k+1]), abs(cor1[j+1,k+1])) / abs(cor0[j+1,k+1] + cor1[j+1,k+1])
    rm <- function(j,k,rho1) (cor0[j+1,k+1] + cor1[j+1,k+1]) / 2 * rho1

    sens2 <- function(k) 2 * min(abs(cor0[1,k+1]), abs(cor1[1,k+1])) / abs(cor0[1,k+1]+ cor1[1,k+1])
    ry <- function(k,rho2) (cor0[1,k+1]+ cor1[1,k+1]) / 2 * rho2
    rho1 <- runif(1, 0, min(apply(expand.grid(1:3,1:3),1, function(x) sens1(x[1],x[2]))))
    rho2 <- runif(1, 0, min(sens2(1),sens2(2), sens2(3)))

    #----- Construct the correlation matrix
    cor.part <- matrix(apply(expand.grid(1:3,1:3), 1, function(x) rm(x[1], x[2], rho1)), 3, 3)
    cor.y1m0 <- matrix(sapply(1:3, function(x) ry(x,rho2)), nrow=1)
    cor.y0m1 <- matrix(sapply(1:3, function(x) ry(x,rho2)), nrow=1)

    COR1 <- 2 * sin(pi * (1/6) *rbind(cbind(cor1,rbind(cor.y1m0, cor.part)), cbind(t(rbind(cor.y1m0, cor.part)), cor0[2:4,2:4] ) ))
    COR0 <- 2 * sin(pi * (1/6) *rbind(cbind(cor0,rbind(cor.y0m1, cor.part)), cbind(t(rbind(cor.y0m1, cor.part)), cor1[2:4,2:4] ) ))

    #----- Check positive definiteness
    S0 <- attr(chol(COR0, pivot=TRUE),"rank");     S1 <- attr(chol(COR1, pivot=TRUE),"rank")
    while (S0 !=7 | S1 !=7){
        rho1 <- runif(1, 0, min(apply(expand.grid(1:3,1:3),1, function(x) sens1(x[1],x[2]))))
        rho2 <- runif(1, 0, min(sens2(1),sens2(2), sens2(3)))
        cor.part <- matrix(apply(expand.grid(1:3,1:3), 1, function(x) rm(x[1], x[2], rho1)), 3, 3)
        cor.y1m0 <- matrix(sapply(1:3, function(x) ry(x,rho2)), nrow=1)
        cor.y0m1 <- matrix(sapply(1:3, function(x) ry(x,rho2)), nrow=1)
        COR1 <- 2 * sin(pi * (1/6) *rbind(cbind(cor1,rbind(cor.y1m0, cor.part)), cbind(t(rbind(cor.y1m0, cor.part)), cor0[2:4,2:4] ) ))
        COR0 <- 2 * sin(pi * (1/6) *rbind(cbind(cor0,rbind(cor.y0m1, cor.part)), cbind(t(rbind(cor.y0m1, cor.part)), cor1[2:4,2:4] ) ))
        S0 <- attr(chol(COR0, pivot=TRUE),"rank");     S1 <- attr(chol(COR1, pivot=TRUE),"rank")
    }

    #----- Index of iterations (on parallel)
    j <- temp

    #----- Index of posterior samples
    index0 <- Thin * j + Burn0; index1 <- Thin * j + Burn1

    #----- Set number of samples (set a large number for smooth surfaces)
    size <- 100000
    s.covariate <- data.frame(X[sample(seq(1,dim(X)[1]), size = size, replace = TRUE),])  #----- Covariates samples
  
  
    #----- Weighting parameters of mixtures
    pi.y0 <- pmax(psi0[index0,], 0)
    pi.y1 <- pmax(psi1[index1,], 0)
    pi.m10 <- pmax(psi1_0[index0,], 0)
    pi.m20 <- pmax(psi2_0[index0,], 0)
    pi.m30 <- pmax(psi3_0[index0,], 0)
    pi.m11 <- pmax(psi1_1[index1,], 0)
    pi.m21 <- pmax(psi2_1[index1,], 0)
    pi.m31 <- pmax(psi3_1[index1,], 0)
  
    ##----- Sampling from the joint distribution of [Y(1;M(1,1,1)), M(0,0,0), M(1,1,1)]
    f1.M0M1 <- function(rho1, k){
        Cor01 <- COR1
      
        #----- Random samples from the Copula and take the standard normal CDF on them
        F <- pnorm(rmnorm(1, mean=rep(0, 7), Cor01), 0, 1)
        X.temp <- s.covariate[k,]
      
        #----- Draw a cluster for each marginal distribution
        clus.y1 <- which(rmultinom(1, 1, pi.y1) == 1)
        clus.m10 <- which(rmultinom(1, 1, pi.m10) == 1)
        clus.m20 <- which(rmultinom(1, 1, pi.m20) == 1)
        clus.m30 <- which(rmultinom(1, 1, pi.m30) == 1)
        clus.m11 <- which(rmultinom(1, 1, pi.m11) == 1)
        clus.m21 <- which(rmultinom(1, 1, pi.m21) == 1)
        clus.m31 <- which(rmultinom(1, 1, pi.m31) == 1)
      
        #----- Samples from marginal distributions
        s.y1 <- qnorm(F[1], mean = gamma10[index1, clus.y1] + gamma11[index1, ] %*% t(X.temp), sd = 1 / sqrt(w1[index1, clus.y1]))
        s.m1_0 <- qnorm(F[5], mean = beta1_00[index0, clus.m10] + beta1_01[index0, ] %*% t(X.temp), sd = 1 / sqrt(s1_0[index0, clus.m10]))
        s.m2_0 <- qnorm(F[6], mean = beta2_00[index0, clus.m20] + beta2_01[index0, ] %*% t(X.temp), sd = 1 / sqrt(s2_0[index0, clus.m20]))
        s.m3_0 <- qnorm(F[7], mean = beta3_00[index0, clus.m30] + beta3_01[index0, ] %*% t(X.temp), sd = 1 / sqrt(s3_0[index0, clus.m30]))
      
        s.m1_1 <- qnorm(F[2], mean = beta1_10[index1, clus.m11] + beta1_11[index1, ] %*% t(X.temp), sd = 1 / sqrt(s1_1[index1, clus.m11]))
        s.m2_1 <- qnorm(F[3], mean = beta2_10[index1, clus.m21] + beta2_11[index1, ] %*% t(X.temp), sd = 1 / sqrt(s2_1[index1, clus.m21]))
        s.m3_1 <- qnorm(F[4], mean = beta3_10[index1, clus.m31] + beta3_11[index1, ] %*% t(X.temp), sd = 1 / sqrt(s3_1[index1, clus.m31]))
        return(c(s.y1, s.m1_0, s.m2_0, s.m3_0, s.m1_1, s.m2_1, s.m3_1))
    }
  
    ##----- Sampling from the joint distribution of [Y(0;M(0,0,0)), M(0,0,0), M(1,1,1)]
    f0.M0M1 <- function(rho1, k){
        Cor01 <- COR0
      
        #----- Random samples from the Copula and take the standard normal CDF on them
        F <- pnorm(rmnorm(1, mean = rep(0, 7), Cor01), 0, 1)
        X.temp <- s.covariate[k, ]
      
        #----- Draw a cluster for each marginal distribution
        clus.y0 <- which(rmultinom(1, 1, pi.y0) == 1)
        clus.m10 <- which(rmultinom(1, 1, pi.m10) == 1)
        clus.m20 <- which(rmultinom(1, 1, pi.m20) == 1)
        clus.m30 <- which(rmultinom(1, 1, pi.m30) == 1)
        clus.m11 <- which(rmultinom(1, 1, pi.m11) == 1)
        clus.m21 <- which(rmultinom(1, 1, pi.m21) == 1)
        clus.m31 <- which(rmultinom(1, 1, pi.m31) == 1)
      
        #----- Samples from marginal distributions
        s.y0 <- qnorm(F[1], mean = gamma00[index0, clus.y0] + gamma01[index0, ] %*% t(X.temp), sd = 1 / sqrt(w0[index0, clus.y0]))
        s.m1_0 <- qnorm(F[2], mean = beta1_00[index0, clus.m10] + beta1_01[index0, ] %*% t(X.temp), sd = 1 / sqrt(s1_0[index0, clus.m10]))
        s.m2_0 <- qnorm(F[3], mean = beta2_00[index0, clus.m20] + beta2_01[index0, ] %*% t(X.temp), sd = 1 / sqrt(s2_0[index0, clus.m20]))
        s.m3_0 <- qnorm(F[4], mean = beta3_00[index0, clus.m30] + beta3_01[index0, ] %*% t(X.temp), sd = 1 / sqrt(s3_0[index0, clus.m30]))
      
        s.m1_1 <- qnorm(F[5], mean = beta1_10[index1, clus.m11] + beta1_11[index1, ] %*% t(X.temp), sd = 1 / sqrt(s1_1[index1, clus.m11]))
        s.m2_1 <- qnorm(F[6], mean = beta2_10[index1, clus.m21] + beta2_11[index1, ] %*% t(X.temp), sd = 1 / sqrt(s2_1[index1, clus.m21]))
        s.m3_1 <- qnorm(F[7], mean = beta3_10[index1, clus.m31] + beta3_11[index1, ] %*% t(X.temp), sd = 1 / sqrt(s3_1[index1, clus.m31]))
        return(c(s.y0, s.m1_0, s.m2_0, s.m3_0, s.m1_1, s.m2_1, s.m3_1))
    }

  
    #----- Estimating Y1 and Y0 conditional on all mediators
    y1.sample <- sapply(seq(1, dim(s.covariate)[1], by=1), function(x) f1.M0M1(rho1, x ))
    y0.sample <- sapply(seq(1, dim(s.covariate)[1], by=1), function(x) f0.M0M1(rho1, x ))
 
    #----- Estimating E[Y1-Y0] given M samples
    y1_1 <- apply(ind1, 1, function(x) mean(y1.sample[1,which(abs(y1.sample[2,]-x[1]) < 0.2 & abs(y1.sample[5,]-x[2]) < 0.2)]))
    y0_1 <- apply(ind1, 1, function(x) mean(y0.sample[1,which(abs(y0.sample[2,]-x[1]) < 0.2 & abs(y0.sample[5,]-x[2]) < 0.2)]))
    d_1 <- y1_1-y0_1
  
    y1_2 <- apply(ind2, 1, function(x) mean(y1.sample[1,which(abs(y1.sample[3,]-x[1]) < 0.2 & abs(y1.sample[6,]-x[2]) < 0.2)]))
    y0_2 <- apply(ind2, 1, function(x) mean(y0.sample[1,which(abs(y0.sample[3,]-x[1]) < 0.2 & abs(y0.sample[6,]-x[2]) < 0.2)]))
    d_2 <- y1_2-y0_2
  
    y1_3 <- apply(ind3, 1, function(x) mean(y1.sample[1,which(abs(y1.sample[4,]-x[1]) < 0.2 & abs(y1.sample[7,]-x[2]) < 0.2)]))
    y0_3 <- apply(ind3, 1, function(x) mean(y0.sample[1,which(abs(y0.sample[4,]-x[1]) < 0.2 & abs(y0.sample[7,]-x[2]) < 0.2)]))
    d_3<-y1_3-y0_3
  
    return(c(d_1, d_2, d_3))
  
}

result<-foreach(temp = 1:n.iter, .combine = rbind) %dopar% main(temp)

save(result, file="SurfaceFrame2.RData")
stopCluster(cl)


